#!/usr/bin/env python3
"""
PyAudio-based VAD detector that uses the same configuration as working record_wav.py
"""

import os
import sys
import time
import threading
from pathlib import Path

# Set optimal device based on PyAudio testing
os.environ['SHERPA_ONNX_MIC_DEVICE'] = '4'  # Device 4 works with PyAudio
os.environ['SHERPA_ONNX_MIC_SAMPLE_RATE'] = '16000'

try:
    import sounddevice as sd
    import numpy as np
    import sherpa_onnx
except ImportError as e:
    print(f"Missing required package: {e}")
    print("Please install with: pip install sounddevice numpy sherpa-onnx")
    sys.exit(1)


class PyAudioVADDetector:
    """
    VAD detector that mimics PyAudio's successful device configuration.
    """
    
    def __init__(self, vad_model_path, device_index=4, sample_rate=16000):
        self.device_index = device_index
        self.sample_rate = sample_rate
        self.is_running = False
        self.speech_detected = False
        
        # Verify model file exists
        if not Path(vad_model_path).is_file():
            raise FileNotFoundError(f"VAD model not found: {vad_model_path}")
        
        # Configure VAD
        config = sherpa_onnx.VadModelConfig()
        config.silero_vad.model = vad_model_path
        config.silero_vad.threshold = 0.5
        config.silero_vad.min_silence_duration = 0.25
        config.silero_vad.min_speech_duration = 0.25
        config.sample_rate = sample_rate
        
        self.vad = sherpa_onnx.VoiceActivityDetector(config, buffer_size_in_seconds=10)
        
        print(f"VAD initialized with device {device_index}")
    
    def start_detection(self):
        """Start VAD detection using the same device configuration as PyAudio."""
        try:
            # Verify the device exists and has input channels
            devices = sd.query_devices()
            if self.device_index >= len(devices):
                print(f"Error: Device {self.device_index} does not exist")
                return False
            
            device_info = devices[self.device_index]
            if device_info['max_input_channels'] == 0:
                print(f"Error: Device {self.device_index} has no input channels")
                return False
            
            print(f"Using device: {device_info['name']}")
            print(f"Input channels: {device_info['max_input_channels']}")
            
            self.is_running = True
            self.detection_thread = threading.Thread(target=self._detection_loop)
            self.detection_thread.daemon = True
            self.detection_thread.start()
            
            print("✅ VAD detection started")
            print("Will print 'true' when speech is detected")
            print("-" * 40)
            return True
            
        except Exception as e:
            print(f"Error starting detection: {e}")
            return False
    
    def stop_detection(self):
        """Stop VAD detection."""
        self.is_running = False
        if hasattr(self, 'detection_thread'):
            self.detection_thread.join(timeout=2.0)
        print("VAD detection stopped")
    
    def _detection_loop(self):
        """Main detection loop using device configuration similar to PyAudio."""
        try:
            # Use the same configuration that works with PyAudio
            with sd.InputStream(
                device=self.device_index,  # Explicit device index
                channels=1,                # Single channel like PyAudio
                samplerate=self.sample_rate,
                dtype='float32',
                blocksize=1024,           # Same chunk size as PyAudio
                latency='low'             # Low latency for real-time
            ) as stream:
                
                print(f"Audio stream opened successfully on device {self.device_index}")
                
                while self.is_running:
                    # Read audio data
                    samples, overflowed = stream.read(1024)
                    
                    if overflowed:
                        print("Warning: Audio buffer overflow")
                    
                    # Ensure correct format
                    if len(samples.shape) > 1:
                        samples = samples[:, 0]  # Take first channel
                    samples = samples.flatten()
                    
                    # Process with VAD
                    self.vad.accept_waveform(samples)
                    
                    # Check for speech detection
                    current_detection = self.vad.is_speech_detected()
                    if current_detection != self.speech_detected:
                        self.speech_detected = current_detection
                        
                        # Print only when speech is detected
                        if self.speech_detected:
                            print("true")
                    
                    # Clean up VAD buffer
                    while not self.vad.empty():
                        self.vad.pop()
        
        except Exception as e:
            print(f"Error in detection loop: {e}")
            self.is_running = False


def main():
    """Main function."""
    model_path = "/home/unitree/softwares/jetson/vad-sherpa-onnx/silero_vad.onnx"
    
    # Allow device selection from command line
    device_index = 4  # Default to device 4
    if len(sys.argv) > 1:
        try:
            device_index = int(sys.argv[1])
        except ValueError:
            print("Usage: python3 vad_pyaudio.py [device_index]")
            sys.exit(1)
    
    try:
        detector = PyAudioVADDetector(model_path, device_index=device_index)
        
        if detector.start_detection():
            try:
                while True:
                    time.sleep(0.1)
            except KeyboardInterrupt:
                print("\nStopping...")
            finally:
                detector.stop_detection()
        else:
            print("Failed to start detection")
            print(f"Try a different device:")
            print(f"python3 record_wav.py find")
            sys.exit(1)
    
    except FileNotFoundError as e:
        print(f"Error: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"Unexpected error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()